import io
import requests
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from threading import Lock, Thread, Event, active_count
import logging
import random

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

url = 'https://dldir1.qq.com/qqfile/qq/QQNT/Windows/QQ_9.9.17_241210_x64_01.exe'  # 替换为你要测试的下载链接

chunk_size = 8192
num_threads = 10  # 并发下载的最大线程数
min_delay = 1  # 最小启动延迟（秒）
max_delay = 5  # 最大启动延迟（秒）
check_interval = 1  # 监控线程检查间隔（秒）
speed_update_interval = 1  # 每秒更新一次速度

progress_lock = Lock()
progresses = {}  # 存储每个线程的进度
total_speeds = {}  # 存储每个线程的总下载速度
downloaded_per_second = {}  # 存储每个线程每秒的下载量
last_downloaded = {}  # 存储每个线程上次更新时的下载量
stop_event = Event()  # 用于优雅地停止线程
total_downloaded = 0  # 全局变量，存储总的下载量

def format_size(size_bytes):
    """将字节数转换为人类可读的格式"""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_bytes < 1024.0:
            return f"{size_bytes:7.2f} {unit}"
        size_bytes /= 1024.0
    return f"{size_bytes:.2f} PB"

def download_file(instance_id):
    start_time = time.time()
    response = requests.get(url, stream=True)
    response.raise_for_status()

    buffer = io.BytesIO()
    total_length = int(response.headers.get('content-length', 0))
    downloaded = 0
    last_update_time = start_time
    last_downloaded[instance_id] = 0

    with progress_lock:
        progresses[instance_id] = {'total': total_length, 'downloaded': 0}
        total_speeds[instance_id] = 0.0
        downloaded_per_second[instance_id] = 0

    for data in response.iter_content(chunk_size=chunk_size):
        if stop_event.is_set():
            break
        if data:
            buffer.write(data)
            downloaded += len(data)

            current_time = time.time()
            elapsed_time = current_time - last_update_time
            if elapsed_time >= speed_update_interval:
                with progress_lock:
                    downloaded_per_second[instance_id] = downloaded - last_downloaded[instance_id]
                    last_downloaded[instance_id] = downloaded
                    progresses[instance_id]['downloaded'] = downloaded
                    global total_downloaded
                    total_downloaded += downloaded_per_second[instance_id]
                last_update_time = current_time

    elapsed_time = time.time() - start_time
    final_speed = total_length / elapsed_time / 1024 / 1024  # 转换为MB/s
    with progress_lock:
        total_speeds[instance_id] = final_speed
        progresses[instance_id]['downloaded'] = total_length
        total_downloaded += total_length - last_downloaded[instance_id]

    buffer.close()  # 关闭字节流以释放内存
    return final_speed

def calculate_and_display_speed():
    while not stop_event.is_set():
        time.sleep(speed_update_interval)  # 等待一秒钟
        with progress_lock:
            total_downloaded_this_second = sum(downloaded_per_second.values())
            total_speed = total_downloaded_this_second / 1024 / 1024  # 转换为MB/s
            formatted_total_downloaded = format_size(total_downloaded)
            print(f'\r当前实时网速: {total_speed:.2f} MB/s | 总下载量: {formatted_total_downloaded}', end='', flush=True)

def run_download_instance(instance_id):
    # 随机启动延迟
    delay = random.uniform(min_delay, max_delay)
    time.sleep(delay)

    while not stop_event.is_set():
        try:
            speed = download_file(instance_id)
            # logging.info(f"实例 {instance_id} 本次下载速度: {speed:.2f} MB/s")
        except requests.exceptions.RequestException as e:
            pass
            # logging.error(f"实例 {instance_id} 下载过程中出错: {e}")
        finally:
            with progress_lock:
                if instance_id in progresses:
                    del progresses[instance_id]  # 清除已完成的下载任务的进度信息
                    del total_speeds[instance_id]  # 清除已完成的下载任务的速度信息
                    del downloaded_per_second[instance_id]  # 清除已完成的下载任务的每秒下载量
                    del last_downloaded[instance_id]  # 清除已完成的下载任务的上次下载量

def monitor_active_threads(executor, num_threads):
    while not stop_event.is_set():
        current_threads = active_count() - 3  # 减去主线程、进度显示线程和速度计算线程
        if current_threads < num_threads:
            new_instance_id = current_threads
            executor.submit(run_download_instance, new_instance_id)
            logging.info(f"启动新实例 {new_instance_id}")
        time.sleep(check_interval)

def main():
    # 创建进度显示和速度计算线程
    speed_thread = Thread(target=calculate_and_display_speed, daemon=True)
    speed_thread.start()

    # 创建线程池
    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        # 启动初始线程
        for i in range(num_threads):
            executor.submit(run_download_instance, i)

        # 创建监控线程
        monitor_thread = Thread(target=monitor_active_threads, args=(executor, num_threads), daemon=True)
        monitor_thread.start()

        try:
            # 等待所有线程完成（实际上由于是无限循环，它们不会自然完成）
            while True:
                time.sleep(1)  # 让主线程空转，等待用户中断
        except KeyboardInterrupt:
            logging.info("\n下载测试已终止。")
            stop_event.set()  # 设置停止事件，通知所有线程停止
            speed_thread.join()  # 等待速度计算线程结束
            monitor_thread.join()  # 等待监控线程结束

if __name__ == "__main__":
    main()